"""
程序功能：
1.加载一个指定路径文件夹内的所有pdf文内容
2.解析所有pdf内容并提取指定内容
3.把解析出来的指定内容写入Excel表格
@author: yczha
@email: yooongchun@foxmail.com
@blog: https://www.yooongchun.com
"""

from cgi import parse_multipart
import os
import re
from typing import List, Tuple
from collections import defaultdict as Ddict

import fitz

from src.utils import get_encoding


class WordFreqCount:
    """关键词词频统计"""

    @staticmethod
    def get_allowed_files(path: str) -> List[str]:
        """读取一个文件夹目录下所有PDF、txt文档路径"""
        file_list = []  # 保存文件地址和名称：name：path
        if not os.path.isdir(path):
            return file_list
        for root, _, files in os.walk(path):
            for file_ in files:
                if file_.endswith((".pdf", ".PDF", ".txt")):  # 判断是否为PDF文件
                    file_list.append(os.path.join(root, file_))
        return file_list

    @staticmethod
    def get_keywords(file_path: str) -> List[str]:
        """关键词"""
        if not os.path.isfile(file_path):
            return []
        encoding = get_encoding(file_path)
        with open(file_path, "r", encoding=encoding) as fin:
            words = fin.readlines()
        keywords = [str(word).strip() for word in words if word and word != ""]
        return keywords

    @staticmethod
    def read_text(filename: str) -> Tuple[bool, str]:
        if filename.endswith(".txt"):
            encoding = get_encoding(filename)
            with open(filename, "r", encoding=encoding) as fin:
                text = fin.read()
                return True, text
        else:
            try:
                doc = fitz.Document(filename)
                text = "".join(page.get_text() for page in doc)
                return True, text
            except Exception as e:
                return False, e

    def run_task(self, file_path: str, keywords: List[str]) -> Tuple[int, str, Ddict]:
        """解析"""
        words_count = Ddict(int)
        state, text = self.read_text(file_path)
        if not state:
            return 1001, text, file_path, words_count
        if len(keywords) < 1:
            return 1002, "没有关键词", file_path, words_count
        for word in keywords:
            words_count[word] += len(list(re.findall(f"{word}", text)))
        return 1000, "", file_path, words_count

    def fold_task(self, params):
        file_path, keywords = params
        return self.run_task(file_path, keywords)


if __name__ == "__main__":
    from src.config import settings

    path: str = os.path.join(settings.work_dir, "example/words_freq_count")
    keywords_file: str = os.path.join(settings.work_dir, "keywords.txt")
    counter = WordFreqCount()
    assert len(counter.get_allowed_files(path)) == len(os.listdir(path))
    assert counter.get_keywords(keywords_file) == ["课程", "贷款", "合同"]
    for filename in counter.get_allowed_files(path):
        code, msg, _, data = counter.run_task(
            filename, counter.get_keywords(keywords_file)
        )
        print(code, msg, data)
